import torch
from dataset import FontDataset
from visualization import show_image_label
import numpy as np
from sklearn.metrics import confusion_matrix
from utils.utils import mask_logits

MODEL_PATH = "checkpoint/mynet-LTH2017-200epochs.pt"
DATA_BASE_URL = "/Users/zenos/Downloads/CCSSD/LTH2017"
NUM_CLASSES = 35

if __name__ == '__main__':
    net = torch.load(MODEL_PATH, map_location='cpu', weights_only=False)
    TestDataset = FontDataset(False, DATA_BASE_URL)
    print(len(TestDataset))
    for i in range(10):
        image, label, classes = TestDataset[i]
        X = image.unsqueeze(0)  # 增加批次维度
        # 创建一个全 0 的张量，形状为 [35]，表示类别总数为 35
        one_hot = torch.zeros(35)
        # 将有效类别对应的索引位置设置为 1
        one_hot[classes] = 1
        # 将其增加批次维度，使得最终形状为 [B, 35]，这里 B=1，表示一个批次的样本
        one_hot = one_hot.unsqueeze(0)
        Y_pred = net(X, one_hot)
        # Y_pred = mask_logits(Y_pred, one_hot)
        Y_pred = Y_pred.argmax(dim=1)

        # 计算混淆矩阵
        # true_labels = label.view(-1).cpu().numpy()
        # pred_labels = Y_pred.view(-1).cpu().numpy()
        # batch_conf_matrix = confusion_matrix(true_labels, pred_labels, labels=np.arange(NUM_CLASSES))
        show_image_label(image.squeeze(0).numpy(), Y_pred.squeeze(0), True)
